/*=========================================================================
 *
 *  Copyright Arash Akbarinia
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/

/**
* @file TrainWithOpenCv.cpp
*
*/

#include "laperconfig.h"
#include "io.h"

#include <ctime>

#include <opencv2/core/core.hpp>
#include <opencv2/ml/ml.hpp>
#include <opencv2/opencv.hpp>

/// classifiers
CvSVM svm;
CvANN_MLP ann_mlp;
CvKNearest knn;
cv::EM em ( cv::EM::DEFAULT_NCLUSTERS, cv::EM::COV_MAT_DIAGONAL, cv::TermCriteria ( CV_TERMCRIT_ITER, ( int ) 1e7, 1e-6 ) );
CvERTrees er_trees;
CvRTrees r_trees;
CvGBTrees gb_trees;
CvBoost boost;
CvDTree d_trees;
CvNormalBayesClassifier nbc;

void ReadAllDescriptors ( const std::vector< std::string >& FileNames, cv::Mat& labels, cv::Mat& descriptors )
{
  for ( unsigned int i = 0; i < FileNames.size(); i++ )
    {
      FILE_LOG ( logINFO ) << "trying to open " << FileNames[i];
      cv::FileStorage fs;
      if ( !fs.open ( FileNames[i], cv::FileStorage::READ ) )
        {
          FILE_LOG ( logWARNING ) << "couldn't open " << FileNames[i];
          continue;
        }

      cv::FileNode PixelDescriptors = fs["PixelDescriptors"];
      cv::FileNodeIterator itb = PixelDescriptors.begin();
      cv::FileNodeIterator ite = PixelDescriptors.end();

      for ( /* */ ; itb != ite; ++itb )
        {
          int label = ( int ) ( *itb ) ["label"];
          labels.push_back ( label );

          std::vector< float > descriptor;
          ( *itb ) ["descriptor"] >> descriptor;
          descriptors.push_back ( cv::Mat ( descriptor ).reshape ( 0, 1 ) );
        }

      fs.release();
    }
}

int main ( int argc, char** argv )
{
  if ( argc < 4 )
    {
      FILE_LOG ( logERROR ) << "usage " << argv[0] << ": InputDirectory OutputDirectory MachineLarningType ";
      return EXIT_FAILURE;
    }

  std::string OutputFolder = argv[2];
  // TODO: check te MlType is suppported
  std::string MlType = argv[3];

  std::vector< std::string > FileNames;
  std::vector< std::string > FileExtensions;
  FileExtensions.push_back ( ".yaml" );
  laper::io::ReadFilesOfDirectory ( argv[1], FileNames, FileExtensions );

  cv::Mat labels;
  cv::Mat descriptors;
  ReadAllDescriptors ( FileNames, labels, descriptors );

  /// ml

  /// svm
  CvSVMParams svm_params;
  svm_params.svm_type    = cv::SVM::C_SVC;
  svm_params.C           = 0.1;
  svm_params.nu          = 0.1;
  svm_params.kernel_type = cv::SVM::LINEAR;
  svm_params.term_crit   = cv::TermCriteria ( CV_TERMCRIT_ITER, ( int ) 1e7, 1e-6 );

  /// knn
//   CvKNearest knn;

  /// ann_mlp
  CvANN_MLP_TrainParams ann_mlp_params;
  ann_mlp_params.term_crit    = cv::TermCriteria ( CV_TERMCRIT_ITER, ( int ) 1e7, 1e-6 );
  ann_mlp_params.train_method = CvANN_MLP_TrainParams::RPROP;

  /// em

  /// r_trees
  int r_trees_min_sample_count               = descriptors.rows * 0.01;
  int r_trees_max_depth                      = 25;
  int r_trees_params_categories              = 10;
  int r_trees_max_num_of_trees_in_the_forest = 100;
  float r_trees_forest_accuracy              = 0.01;
  int r_trees_termcrit_type                  = CV_TERMCRIT_ITER | CV_TERMCRIT_EPS;
  CvRTParams r_trees_params ( r_trees_max_depth, r_trees_min_sample_count, 0.0, false, r_trees_params_categories, 0, false, 0, r_trees_max_num_of_trees_in_the_forest, r_trees_forest_accuracy, r_trees_termcrit_type );

  /// er_trees
  CvRTParams er_trees_params ( r_trees_max_depth, r_trees_min_sample_count, 0.0, false, r_trees_params_categories, 0, false, 0, r_trees_max_num_of_trees_in_the_forest, r_trees_forest_accuracy, r_trees_termcrit_type );

  /// gb_trees
  int gb_trees_params_max_depth    = 3;
  int gb_trees_params_weak_count   = 600;
  float gb_trees_shrinkage         = 0.01;
  float gb_trees_subsample_portion = 0.8;
  CvGBTreesParams gb_trees_params ( CvGBTrees::DEVIANCE_LOSS, gb_trees_params_weak_count, gb_trees_shrinkage, gb_trees_subsample_portion, gb_trees_params_max_depth, false );

  /// boost
  int boost_weak_count          = 600;
  double boost_weight_trim_rate = 0.99;
  int boost_max_depth           = 1;
  CvBoostParams boost_params ( CvBoost::REAL, boost_weak_count, boost_weight_trim_rate, boost_max_depth, true, 0 );
  boost_params.min_sample_count = descriptors.rows * 0.01;

  /// d_trees
  CvDTreeParams d_trees_params;
  d_trees_params.min_sample_count = descriptors.rows * 0.01;
  d_trees_params.max_depth        = 25;
  d_trees_params.cv_folds         = 10;
  d_trees_params.max_categories   = 10;

  /// nbc

  // write responsed to file
  std::string MlFileName;
  MlFileName.append ( OutputFolder );
  MlFileName.append ( "ml_" );
  MlFileName.append ( MlType );
  MlFileName.append ( ".yaml" );

  FILE_LOG ( logDEBUG ) << "starting to train " << MlType;
  time_t start;
  time_t end;
  time ( &start );
  if ( MlType.compare ( "svm" ) == 0 )
    {
      //   svm.train_auto ( descriptors, labels, Mat(), Mat(), svm_params, 10 );
      svm.train ( descriptors, labels, cv::Mat(), cv::Mat(), svm_params );
      svm.save ( MlFileName.c_str ( ) );
    }
  else if ( MlType.compare ( "knn" ) == 0 )
    {
      knn.train ( descriptors, labels );
      knn.save ( MlFileName.c_str ( ) );
    }
  else if ( MlType.compare ( "ann_mlp" ) == 0 )
    {
      ann_mlp.train ( descriptors, labels, cv::Mat(), cv::Mat(), ann_mlp_params );
      ann_mlp.save ( MlFileName.c_str ( ) );
    }
  else if ( MlType.compare ( "em" ) == 0 )
    {
      em.train ( descriptors );
      // FIXME: can you save it?
//       em.write ( ml_filename.c_str ( ) );
    }
  else if ( MlType.compare ( "er_trees" ) == 0 )
    {
      er_trees.train ( descriptors, CV_ROW_SAMPLE, labels, cv::Mat(), cv::Mat(), cv::Mat(), cv::Mat(), er_trees_params );
      er_trees.save ( MlFileName.c_str ( ) );
    }
  else if ( MlType.compare ( "r_trees" ) == 0 )
    {
      r_trees.train ( descriptors, CV_ROW_SAMPLE, labels, cv::Mat(), cv::Mat(), cv::Mat(), cv::Mat(), r_trees_params );
      r_trees.save ( MlFileName.c_str ( ) );
    }
  else if ( MlType.compare ( "gb_trees" ) == 0 )
    {
      gb_trees.train ( descriptors, CV_ROW_SAMPLE, labels, cv::Mat(), cv::Mat(), cv::Mat(), cv::Mat(), gb_trees_params );
      gb_trees.save ( MlFileName.c_str ( ) );
    }
  else if ( MlType.compare ( "boost" ) == 0 )
    {
      boost.train ( descriptors, CV_ROW_SAMPLE, labels, cv::Mat(), cv::Mat(), cv::Mat(), cv::Mat(), boost_params );
      boost.save ( MlFileName.c_str ( ) );
    }
  else if ( MlType.compare ( "d_trees" ) == 0 )
    {
      d_trees.train ( descriptors, CV_ROW_SAMPLE, labels, cv::Mat(), cv::Mat(), cv::Mat(), cv::Mat(), d_trees_params );
      d_trees.save ( MlFileName.c_str ( ) );
    }
  else if ( MlType.compare ( "nbc" ) == 0 )
    {
      nbc.train ( descriptors, labels );
      nbc.save ( MlFileName.c_str ( ) );
    }

  time ( &end );
  FILE_LOG ( logDEBUG ) << "the training is finished";
  FILE_LOG ( logDEBUG ) << "elapssed time is " << difftime ( end, start );

  return 0;
}